package com.yahoo.labs.cluster;

import it.unimi.dsi.io.FileLinesCollection;
import it.unimi.dsi.io.FileLinesCollection.FileLinesIterator;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.nio.charset.Charset;

import cern.colt.matrix.impl.SparseDoubleMatrix2D;

import com.yahoo.labs.snow.util.MatrixUtils;

/**
 * A wrapper around METIS.
 * 
 * To use this class you need to install the <a
 * href="http://glaros.dtc.umn.edu/gkhome/views/metis">METIS package</a> and make sure the programs
 * (particularly {@link #METIS_COMMAND}) are in your execution PATH.
 * 
 * @author chato
 * 
 */
public class MetisWrapper {

	/**
	 * Metis deals with integer weights. Weights are multiplied by this factor and rounded before
	 * passing them.
	 */
	private static final int PRECISION_DOUBLE_TO_INTEGER = 1000;

	final static String METIS_COMMAND = "kmetis";

	MetisWrapper() {
		String output = run(new String[] { METIS_COMMAND });
		if (!output.startsWith("Usage:")) {
			throw new IllegalStateException("Expected output of " + METIS_COMMAND + " to start with 'Usage:'");
		}
	}

	static String run(String[] command) {
		ProcessBuilder pb = new ProcessBuilder(command);
		pb.redirectErrorStream(true);
		Process process;
		try {
			process = pb.start();
		} catch (IOException e) {
			System.err.println("Can't find program in your PATH. Please make sure '" + METIS_COMMAND + "' is in your PATH");
			throw new IllegalStateException(METIS_COMMAND + ": " + e);
		}
		BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
		String line;
		StringBuffer output = new StringBuffer();
		try {
			while ((line = reader.readLine()) != null) {
				output.append(line);
				output.append("\n");
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
		return output.toString();
	}

	int[] cluster(SparseDoubleMatrix2D graph, int nClusters) {
		if (graph.rows() != graph.columns()) {
			throw new IllegalArgumentException("Input should be a square matrix");
		}
		int nodes = graph.rows();
		int edges = MatrixUtils.countNonZeros(graph);

		File inputFile;
		try {
			inputFile = File.createTempFile("metis-input-", ".tmp");
		} catch (IOException e) {
			throw new IllegalStateException("Could not create tmp file:" + e);
		}

		PrintWriter pw;
		try {
			pw = new PrintWriter(new FileOutputStream(inputFile));
		} catch (FileNotFoundException e) {
			throw new IllegalStateException("Could not find created tmp file:" + e);
		}
		pw.println(nodes + " " + (int) (edges / 2) + " " + 1); // 1=weighted graph
		for (int node = 0; node < nodes; node++) {
			boolean firstNeighbor = true;
			for (int neighbor = 0; neighbor < nodes; neighbor++) {
				if (graph.getQuick(node, neighbor) > 0) {

					if (graph.getQuick(neighbor, node) != graph.getQuick(node, neighbor)) {
						throw new IllegalArgumentException("Input matrix of weights is not symmetric (" + node + "," + neighbor + ")=" + graph.getQuick(node, neighbor)
								+ " but " + "(" + neighbor + "," + node + ")=" + graph.getQuick(neighbor, node));
					}

					if (firstNeighbor) {
						firstNeighbor = false;
					} else {
						pw.print(" ");
					}
					pw.print((neighbor + 1) + " " + (int) Math.round(PRECISION_DOUBLE_TO_INTEGER * graph.getQuick(node, neighbor)));
				}
			}
			pw.println();
		}
		pw.close();

		String output = run(new String[] { METIS_COMMAND, inputFile.getAbsoluteFile().toString(), Integer.toString(nClusters) });

		File outputFile = new File(inputFile.getAbsoluteFile().toString() + ".part." + Integer.toString(nClusters));
		if (!outputFile.exists()) {
			System.err.println(output);
			System.err.println("Input file: " + inputFile.getAbsoluteFile().toString());
			throw new IllegalStateException("There is no output file " + outputFile);
		}

		FileLinesCollection flc = new FileLinesCollection(outputFile.getAbsoluteFile().toString(), Charset.forName("UTF-8").name());
		FileLinesIterator it = flc.iterator();
		int[] result = new int[nodes];
		for (int node = 0; node < nodes; node++) {
			int cluster = Integer.parseInt(it.next().toString());
			result[node] = cluster;
		}

		return result;
	}
}
